{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torch import nn\n",
    "import argparse\n",
    "import torch.optim\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean = [0.5070751592371323, 0.48654887331495095, 0.4409178433670343]\n",
    "std = [0.2673342858792401, 0.2564384629170883, 0.27615047132568404]\n",
    "num_workers = 2\n",
    "\n",
    "\n",
    "def cifar100_dataset(args):\n",
    "    transform_train = transforms.Compose([\n",
    "        transforms.RandomCrop(32, padding=4),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.RandomRotation(15),  # 数据增强\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean, std)\n",
    "    ])\n",
    "    transform_test = transforms.Compose(\n",
    "        [transforms.ToTensor(),\n",
    "         transforms.Normalize(mean, std)])\n",
    "\n",
    "    cifar100_training = torchvision.datasets.CIFAR100(root=args.data_path, train=True, download=True, transform=transform_train)\n",
    "    train_loader = torch.utils.data.DataLoader(cifar100_training, batch_size=args.batch_size, shuffle=True, num_workers=num_workers)\n",
    "\n",
    "    cifar100_testing = torchvision.datasets.CIFAR100(root=args.data_path, train=False, download=True, transform=transform_test)\n",
    "    test_loader = torch.utils.data.DataLoader(cifar100_testing, batch_size=100, shuffle=False, num_workers=num_workers)\n",
    "\n",
    "    return train_loader, test_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SE_Block(nn.Module):                         # Squeeze-and-Excitation block\n",
    "    def __init__(self, in_planes):\n",
    "        super(SE_Block, self).__init__()\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.conv1 = nn.Conv2d(in_planes, in_planes // 16, kernel_size=1)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.conv2 = nn.Conv2d(in_planes // 16, in_planes, kernel_size=1)\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.avgpool(x)\n",
    "        x = self.conv1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        out = self.sigmoid(x)\n",
    "        return out\n",
    "\n",
    "\n",
    "class BasicBlock(nn.Module):      # 左侧的 residual block 结构（18-layer、34-layer）\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, in_planes, planes, stride=1):      # 两层卷积 Conv2d + Shutcuts\n",
    "        super(BasicBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3,\n",
    "                               stride=stride, padding=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(planes)\n",
    "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,\n",
    "                               stride=1, padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "\n",
    "        self.SE = SE_Block(planes)           # Squeeze-and-Excitation block\n",
    "\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_planes != self.expansion*planes:      # Shutcuts用于构建 Conv Block 和 Identity Block\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_planes, self.expansion*planes,\n",
    "                          kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(self.expansion*planes)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        out = self.bn2(self.conv2(out))\n",
    "        SE_out = self.SE(out)\n",
    "        out = out * SE_out\n",
    "        out += self.shortcut(x)\n",
    "        out = F.relu(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Bottleneck(nn.Module):      # 右侧的 residual block 结构（50-layer、101-layer、152-layer）\n",
    "    expansion = 4\n",
    "\n",
    "    def __init__(self, in_planes, planes, stride=1):      # 三层卷积 Conv2d + Shutcuts\n",
    "        super(Bottleneck, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(planes)\n",
    "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,\n",
    "                               stride=stride, padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "        self.conv3 = nn.Conv2d(planes, self.expansion*planes,\n",
    "                               kernel_size=1, bias=False)\n",
    "        self.bn3 = nn.BatchNorm2d(self.expansion*planes)\n",
    "\n",
    "        self.SE = SE_Block(self.expansion*planes)           # Squeeze-and-Excitation block\n",
    "\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_planes != self.expansion*planes:      # Shutcuts用于构建 Conv Block 和 Identity Block\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_planes, self.expansion*planes,\n",
    "                          kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(self.expansion*planes)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        out = F.relu(self.bn2(self.conv2(out)))\n",
    "        out = self.bn3(self.conv3(out))\n",
    "        SE_out = self.SE(out)\n",
    "        out = out * SE_out\n",
    "        out += self.shortcut(x)\n",
    "        out = F.relu(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class SE_ResNet(nn.Module):\n",
    "    def __init__(self, block, num_blocks, num_classes=1000):\n",
    "        super(SE_ResNet, self).__init__()\n",
    "        self.in_planes = 64\n",
    "\n",
    "        self.conv1 = nn.Conv2d(3, 64, kernel_size=3,\n",
    "                               stride=1, padding=1, bias=False)                  # conv1\n",
    "        self.bn1 = nn.BatchNorm2d(64)\n",
    "        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)       # conv2_x\n",
    "        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)      # conv3_x\n",
    "        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)      # conv4_x\n",
    "        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)      # conv5_x\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.linear = nn.Linear(512 * block.expansion, num_classes)\n",
    "\n",
    "    def _make_layer(self, block, planes, num_blocks, stride):\n",
    "        strides = [stride] + [1]*(num_blocks-1)\n",
    "        layers = []\n",
    "        for stride in strides:\n",
    "            layers.append(block(self.in_planes, planes, stride))\n",
    "            self.in_planes = planes * block.expansion\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.bn1(self.conv1(x)))\n",
    "        x = self.layer1(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.layer3(x)\n",
    "        x = self.layer4(x)\n",
    "        x = self.avgpool(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        out = self.linear(x)\n",
    "        return out\n",
    "\n",
    "\n",
    "def SE_ResNet18():\n",
    "    return SE_ResNet(BasicBlock, [2, 2, 2, 2])\n",
    "\n",
    "\n",
    "def SE_ResNet34():\n",
    "    return SE_ResNet(BasicBlock, [3, 4, 6, 3])\n",
    "\n",
    "\n",
    "def SE_ResNet50():\n",
    "    return SE_ResNet(Bottleneck, [3, 4, 6, 3])\n",
    "\n",
    "\n",
    "def SE_ResNet101():\n",
    "    return SE_ResNet(Bottleneck, [3, 4, 23, 3])\n",
    "\n",
    "\n",
    "def SE_ResNet152():\n",
    "    return SE_ResNet(Bottleneck, [3, 8, 36, 3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(net, train_loader, optimizer, criterion, writer, args, epoch, index_num):\n",
    "    net.train()\n",
    "    train_tqdm = tqdm(train_loader, desc=\"Epoch \" + str(epoch))\n",
    "    for index, (inputs, labels) in enumerate(train_tqdm):\n",
    "        optimizer.zero_grad()\n",
    "        outputs = net(inputs.to(args.device))\n",
    "        loss = criterion(outputs, labels.to(args.device))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        writer.add_scalar(\"loss/train\", loss, index_num)\n",
    "        index_num = index_num + 1\n",
    "        train_tqdm.set_postfix({\"loss\": \"%.3g\" % loss.item()})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate(net, test_loader, criterion, writer, args, epoch, loss_vector, accuracy_vector):\n",
    "    net.eval()\n",
    "    val_loss, correct = 0, 0\n",
    "    for index, (data, target) in enumerate(test_loader):\n",
    "        data = data.to(args.device)\n",
    "        target = target.to(args.device)\n",
    "        output = net(data)\n",
    "        val_loss += criterion(output, target.to(args.device)).data.item()\n",
    "        pred = output.data.max(1)[1]\n",
    "        correct += pred.eq(target.data).cpu().sum()\n",
    "    val_loss /= len(test_loader)\n",
    "    loss_vector.append(val_loss)\n",
    "    writer.add_scalar(\"loss/validation\", val_loss, epoch)\n",
    "    accuracy = 100. * correct.to(torch.float32) / len(test_loader.dataset)\n",
    "    accuracy_vector.append(accuracy)\n",
    "    writer.add_scalar(\"accuracy/validation\", accuracy, epoch)\n",
    "\n",
    "    print(\"***** Eval results *****\")\n",
    "    print('epoch: {}, Validation set: Average loss: {:.4f}, Accuracy: {}/{} ({:.4f}%)\\n'.format(\n",
    "        epoch, val_loss, correct, len(test_loader.dataset), accuracy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0: 100%|██████████| 391/391 [00:20<00:00, 18.87it/s, loss=3.3] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "***** Eval results *****\n",
      "epoch: 0, Validation set: Average loss: 3.2297, Accuracy: 2141/10000 (21.4100%)\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAUIAAADSCAYAAAA/vMlrAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8qNh9FAAAACXBIWXMAAAsTAAALEwEAmpwYAAARZUlEQVR4nO3dfZBV9X3H8fdHnrSiFYVmlEefyPiUolmpJtU4xIAxBjVplRgn6KSxNhrbSZ1JmDQhYtsYbdJq44yaqU1jBqnSmiFVq/hArFHURREFY3kQFXwABZIwGKLw7R/nt/Z4vXf37u69ey/8Pq+ZO5x7nvb729XPnnPPnvNVRGBmlrM9Wl2AmVmrOQjNLHsOQjPLnoPQzLLnIDSz7DkIzSx7DkJrKEmnSFpXer9c0in1rNuHr3WDpG/2dftu9vttST9p9H6tfQ1udQG2e4uIoxqxH0kXAH8WEX9c2vfFjdi3mY8IzSx7DkJ7H0lfkzS/Yt61kq5L0xdKek7SbyStkfTn3exrraRT0/Rekn4kabOkFcDxFet+XdLqtN8Vks5O848AbgBOlLRV0pY0/0eS/ra0/ZckrZK0SdICSQeVloWkiyWtlLRF0vWSVOf3Y3o6xd8iaVGqp/y9Wp9qfl7Sx9P8yZI6Jf1a0uuSvl/P17LWcBBaNfOA0yXtAyBpEHAOMDct3wCcAewLXAj8o6Tj6tjvbODQ9JoGzKxYvho4Cfh94ArgJ5IOjIjngIuBRyNieETsV7ljSVOA76Q6DwReTOMoO4MifD+U1pvWU8GSJgK3An8FjALuAn4maaikDwKXAsdHxD5pf2vTptcC10bEvmm8t/X0tax1HIT2PhHxIvAkcHaaNQXYFhGL0/I7I2J1FH4O3EsRYD05B/i7iNgUES8D11V83dsj4pWI2BkR/w6sBCbXWfbngZsj4smI2A7MojiCnFBa56qI2BIRLwEPApPq2O+5wJ0RsTAi3gb+AdgL+AiwAxgGHClpSESsjYjVabu3gcMkjYyIrV3fO2tPDkKrZS7wuTR9Hv9/NIikT0panE5BtwCnAyPr2OdBwMul9y+WF0r6gqSl6RR0C3B0nfvt2ve7+4uIrcCbwOjSOq+VprcBw/uw351pDKMjYhXFkeK3gQ2S5pVOx78ITAR+KekJSWfUOQ5rAQeh1XI7cIqkMRRHhnMBJA0D/oPiyOgD6TT1LqCez9teBcaW3o/rmpA0HvghxanmAWm/z5b229Njkl4Bxpf2tzdwALC+jrp6s19RjGE9QETMTVeyx6cav5vmr4yIzwF/kObNTzVZG3IQWlURsRFYBPwr8EL6nA5gKMXp4EbgHUmfBKbWudvbgFmSRqSA/Upp2d4UQbIRigsyFEeEXV4HxkgaWmPftwIXSpqUwvrvgcciYm2dtXVX86ckfVzSEOCvge3AI5I+KGlK+nq/Bd4Cdqb6z5c0Kh1Bbkn72tnPWqxJHITWnbnAqZROiyPiN8BlFAGxmeK0eUGd+7uC4jTzBYrPFW8p7XcF8D3gUYrQOwb4RWnbB4DlwGuS3qjccUTcB3yT4mj1VYoLFDPqrKumiHgeOB/4Z+AN4NPApyPidxS/EK5K81+jOPqblTY9DVguaSvFhZMZEfFWf+ux5pAfzGpmufMRoZllz0FoZtlzEJpZ9hyEZpY9B6GZZa/Hx3BJ2hN4iOJPBQYD8yNido11PwvMp7j3sjPNm0XxV/Y7gMsi4p7uvt7IkSNjwoQJvRmDmVmPlixZ8kZEjKq2rJ7nEW4HpkTE1vQHpQ9Lurvy3sl0g/5fAo+V5h1J8bdcR1HcqnSfpIkRsaPWF5swYQKdnZ11lGVmVj9JL9Za1uOpcbqxfmt6OyS9qv3x4ZUUtxL9tjTvTGBeRGyPiBeAVdR/E72Z2YCo6zNCSYMkLaV4/NLCiHisYvlxwNiIuLNi09G89yb7dbz3Jviu7S9Kz27r3LhxY2/qNzPrt7qCMCJ2RMQkYAwwWdK794BK2gP4PsU9mH0SETdFREdEdIwaVfUU3sysaXp11TgitlA8x+200ux9KG6OXyRpLXACsEBSB8UTOspPGxlD/58GYmbWUD0GoaRRkvZL03sBnwB+2bU8In4VESMjYkJETAAWA9PTVeMFwAxJwyQdDBwOPN74YZiZ9V09V40PBP4tPa59D+C2iPgvSXOAzoio+eSRiFgu6TZgBfAOcEl3V4zNzFqh7Z4+09HREf7zGTNrNElLIqKj2jLfWWJm2XMQmln2HIRmlj0HoZllz0FoZtlzEJpZ9hyEZpY9B6GZZc9BaGbZcxCaWfYchGaWPQehmWXPQWhm2XMQmln26nkw656SHpf0tKTlkq6oss7Fkp6RtFTSw6l7HZImSHorzV8q6YZmDMLMrD8a1c5zbkTcACBpOkUPk67H+a9O/U7MzNpSj0EYxZNbu23nGRG/Lr3du3K5mVk7a0g7z7TOJZJWA1cDl5UWHSzpKUk/l3RSjf27naeZtUyvHtWfmjjdAXwlIp6tsc55wLSImClpGDA8It6U9GHgp8BRFUeQ7+FH9ZtZMzTsUf012nlWmgecldbfHhFvpuklwGpgYm++pplZs/W7nWeaf3jp7aeAlaVtB6XpQyjaea5pSOVmZg3SqHael0o6FXgb2AzMTNueDMyR9DawE7g4IjY1fBRmZv3gdp5mlgW38zQz64aD0Myy5yA0s+w5CM0sew5CM8ueg9DMsucgNLPsOQjNLHsOQjPLnoPQzLLnIDSz7DkIzSx7DkIzy15Tu9ilZbMkrZL0vKRpjR6AmVl/NbWLXQrEGcBRwEHAfZImRsSOxg7DzKzvejwijEJfu9idCcxLj+x/AVgFTO531WZmDVTPESHp6dRLgMOA62t1sQO+CgwFpqTZo4HykeO6NM/MrG3UdbEkInakJu1jgMmSjq6yzvURcSjwNeBvelOE23maWSs1tYsdsB4YW1o2Js2r3O9NEdERER2jRo3qTUlmZv3W1C52wAJghqRhkg6m6GL3eAPqNjNrmKZ2sYuI5ZJuA1YA7wCX+IqxmbUbd7Ezsyy4i52ZWTcchGaWPQehmWXPQWhm2XMQmln2HIRmlj0HoZllz0FoZtlzEJpZ9hyEZpY9B6GZZc9BaGbZcxCaWfYchGaWvUa18/yqpBWSlkm6X9L40rIdqc3nUkkLGj0AM7P+alQ7z6eAjojYJukvgKuBc9Oyt1K/EzOzttSodp4PRsS29HYxRW8SM7NdQl2fEUoaJGkpsAFYWK2dZ8kXgbtL7/dMHeoWSzqrz5WamTVJXX2NU5+RSamJ0x2Sjo6IZyvXk3Q+0AF8rDR7fESsl3QI8ICkZyJidcV2FwEXAYwbN65vIzEz66OGtfNMzZu+AUyPiO2lbdanf9cAi4Bjq+zX7TzNrGUa1c7zWOBGihDcUJo/QtKwND0S+ChFRzszs7bRqHae1wDDgdslAbwUEdOBI4AbJe1M214VEQ5CM2srPQZhRCyj+unst0rTp9bY9hHgmP4UaGbWbL6zxMyy5yA0s+w5CM0sew5CM8ueg9DMsucgNLPsOQjNLHsOQjPLnoPQzLLnIDSz7DkIzSx7DkIzy56D0Myy5yA0s+wNRDvPmZJWptfMRg/AzKy/6jki7Grn+YfAJOA0SSdUrNPVzvNDwHyKdp5I2h+YDfwRMBmYLWlEg2o3M2uIZrfznEbR9W5TRGwGFlKl34mZWSs1u53naODl0rJ1aZ6ZWduoKwgjYkdETKI40pss6ehq65XaeV7TmyIkXZR6H3du3LixN5uamfVbs9t5rgfGllYbk+ZV7tftPM2sZZrazhO4B5ia2nqOAKameWZmbaOp7TwjYpOkK4En0r7mRMSmxg/DzKzvFBE9rzWAOjo6orOzs9VlmNluRtKSiOiotsx3lphZ9hyEZpY9B6GZZc9BaGbZcxCaWfYchGaWPQehmWXPQWhm2XMQmln2HIRmlj0HoZllz0FoZtlzEJpZ9hyEZpa9RrXzPFnSk5LekfQnFct2SFqaXgsaWbyZWSPU82DWrnaeWyUNAR6WdHdELC6t8xJwAXB5le3fSv1OzMzaUo9BGMWTW3tq57kWQNLOBtdnZtZ0zWjnWWnP1KFusaSz+lCjmVlTNbSdZw3j0+OxzwP+SdKhlSu4naeZtVLD2nl2s8369O8aYBFwbJV13M7TzFqmIe08u9l2hKRhaXok8FFgRZ+rNTNrgnqOCA8EHpS0jKIt58Kudp6SpgNIOl7SOuBPgRslLU/bHgF0Snqa4kjyqohwEJpZW6nnqvEyqp/Ofqs0/QTF54eV6zwCHNPPGs3Mmsp3lphZ9hyEZpY9B6GZZc9BaGbZcxCaWfYchGaWPQehmWXPQWhm2XMQmln2HIRmlj0HoZllz0FoZtlzEJpZ9hyEZpa9gWjnOVPSyvSa2cjizcwaoantPCXtD8wGOig63y2RtCAiNjekejOzBujxiDAKPbbzTA9wrWznOY3iidabUvgtpBf9TszMBkKz23mOBl4uvV+X5pmZtY2BaOfZI7fzNLNWanY7z/XA2NL7MWle5X7dztPMWqap7TyBe4Cpqa3nCGBqmmdm1jaa2s4zIjYBV6btngDmpHlmZm1DEdHzWgOoo6MjOjs7W12Gme1mJC2JiI5qy3xniZllz0FoZtlzEJpZ9hyEZpY9B6GZZc9BaGbZcxCaWfYchGaWPQehmWXPQWhm2XMQmln2HIRmlj0HoZllz0FoZtlzEJpZ9hyEZpY9B6GZZa/tnlAtaSPwYqvrqGIk8Eari2gQj6U97S5jaddxjI+Iqt3h2i4I25WkzlqP+d7VeCztaXcZy644Dp8am1n2HIRmlj0HYf1uanUBDeSxtKfdZSy73Dj8GaGZZc9HhGaWPQchIOk0Sc9LWiXp61WWj5d0v6RlkhZJGlNaNk7SvZKek7RC0oQBLb5CP8dytaTlaSzXSdLAVv+eOm+WtEHSszWWK9W4Ko3luNKymZJWptfMgau6ur6ORdIkSY+mn8kySecObOXvq7PPP5O0fF9J6yT9YGAq7oWIyPoFDAJWA4cAQ4GngSMr1rkdmJmmpwC3lJYtAj6RpocDv7crjgX4CPCLtI9BwKPAKS0cy8nAccCzNZafDtwNCDgBeCzN3x9Yk/4dkaZHtPi/sb6OZSJweJo+CHgV2G9XG0dp+bXAXOAHrfx5VHv5iBAmA6siYk1E/A6YB5xZsc6RwANp+sGu5ZKOBAZHxEKAiNgaEdsGpuyq+jwWIIA9KQJ0GDAEeL3pFdcQEQ8Bm7pZ5Uzgx1FYDOwn6UBgGrAwIjZFxGZgIXBa8yuura9jiYj/jYiVaR+vABuAqn8QPBD68TNB0oeBDwD3Nr/S3nMQwmjg5dL7dWle2dPAZ9L02cA+kg6g+I29RdJ/SnpK0jWSBjW94tr6PJaIeJQiGF9Nr3si4rkm19sftcZaz/eg3fRYs6TJFL+kVg9gXb1VdRyS9gC+B1zekqrq4CCsz+XAxyQ9BXwMWA/sAAYDJ6Xlx1Ockl7QohrrVXUskg4DjgDGUPwHPUXSSa0r07qko6pbgAsjYmer6+mDLwN3RcS6VhdSy+BWF9AG1gNjS+/HpHnvSqclnwGQNBz4bERskbQOWBoRa9Kyn1J8NvIvA1B3Nf0Zy5eAxRGxNS27GzgR+J+BKLwPao11PXBKxfxFA1ZV39T8uUnaF7gT+EY63WxntcZxInCSpC9TfI4+VNLWiHjfxbxW8REhPAEcLulgSUOBGcCC8gqSRqbDe4BZwM2lbfeT1PW5zRRgxQDUXEt/xvISxZHiYElDKI4W2/nUeAHwhXSl8gTgVxHxKnAPMFXSCEkjgKlpXjurOpb0M7yD4nO3+a0tsS5VxxERn4+IcRExgeKM5MftFILgI0Ii4h1Jl1L8zzIIuDkilkuaA3RGxAKKI4zvSArgIeCStO0OSZcD96c/NVkC/LAV40j19HkswHyKIH+G4sLJf0fEzwZ6DF0k3UpR68h05D2b4gIOEXEDcBfFVcpVwDbgwrRsk6QrKX4pAMyJiO4+4G+6vo4FOIfiSu0Bki5I8y6IiKUDVXtZP8bR9nxniZllz6fGZpY9B6GZZc9BaGbZcxCaWfYchGaWPQehmWXPQWhm2XMQmln2/g/C7u/X9hL5iwAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 360x216 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAUIAAADSCAYAAAA/vMlrAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8qNh9FAAAACXBIWXMAAAsTAAALEwEAmpwYAAASpUlEQVR4nO3de5RdZX3G8e9jLqgklEjGlNwYL3GVURDoGKEBxYgYsqwIXSooSGBpaoWWtFALWpfXoksqVheuUjTUW5AiJDaVAImYLIxCdIIBSQYkRDCXAQYQk5R20eCvf+x37O7hnMyZOWfmHOZ9PmudNfvsd7/v+b1zeWZf5uxRRGBmlrMXtLoAM7NWcxCaWfYchGaWPQehmWXPQWhm2XMQmln2HIT2/0g6UdKO0vPNkk6sZ9thvNZVkj423P5mzTK+1QVYe4uIVzdjHEmLgPdHxPGlsT/YjLHNGuU9QrMmk+QdjOcZB+EYJOnvJN1Qse5Lkr6cls+V1Ctpj6Rtkv58P2M9JOmktPwiSV+X9BtJW4DXVWx7iaQH07hbJJ2W1h8OXAUcJ2mvpKfS+q9L+kyp/wckbZX0pKSVkqaX2kLSByU9IOkpSV+RpBo1z5V0R9quT9KVkiaW2l8taU16nUclfSStHyfpI6U5bJQ0S1Jnev3xpTHWSXp/Wl4k6ceSvijpCeATkl4h6YeSnpD0uKRlkg4u9Z8labmk/rTNlZImppqOKG33UklPS+qo9TWyxjkIx6brgIWSJkPxAw68C7g2tT8GvA04CDgX+KKkY+oY9+PAK9LjrcA5Fe0PAicAfwB8Evi2pEMjohf4IHBHREyKiIMrB5Y0H/hsqvNQ4OE0j7K3UYTvkWm7t9ao81ngr4GpwHHAm4EPpdeZDPwAuAWYDrwSuC31+xvgTGAhxefmPODp/Xw+yl4PbAOmAf8AKM1nOnA4MAv4RKphHPD9NMdOYAZwXUQ8k+Z8VmncM4HbIqK/zjpsOCLCjzH4ANYD70vLbwEe3M+23wMuTMsnAjtKbQ8BJ6XlbcCCUtvi8rZVxt0EnJqWFwHrK9q/DnwmLS8FPl9qmwT8D9CZngdwfKn9euCSOj8XS4AVaflM4Oc1trt/oN6K9Z3p9ceX1q2jOOc5MLdfD1LDOwZelyKc+8vjlbZ7PfBrQOl5D/CuVn8/jfWH9wjHrmspfugB3sP/7Q0i6RRJd6bDsKco9oCm1jHmdGB76fnD5UZJ75O0KR2SPgW8ps5xB8b+/XgRsRd4gmJvacAjpeWnKcLyOSS9StL3JT0iaTdwWamOWRR7rtXsr20w5c8LkqZJuk7SzlTDtytqeDgi9lUOEhEbKOZ2oqQ/othjXTnMmqxODsKx67sUP0wzgdNIQSjpAOBG4B+BaVEcpq6iOJQbTB/FD/GA2QMLkg4DvgpcABySxr23NO5gtznaBRxWGu9A4BBgZx11Vfpn4D5gTkQcBHykVMd24OU1+m2nOOyv9J/p44tL6/6wYpvK+V2W1h2RajiroobZ+7mo8o20/dnADRHx3zW2syZxEI5RUZxTWgf8K/CrKM7TAUwEDqA4NNsn6RTg5DqHvR64VNKUFLB/WWo7kOIHvx+KCzIUe4QDHgVmli9aVPgOcK6ko1JYXwZsiIiH6qytbDKwG9ib9qr+otT2feBQSUskHSBpsqTXp7avAZ+WNEeFIyUdkj6XO4Gz0gWV86gemJU17AV+K2kG8Leltp9S/FL5nKQDJb1Q0rxS+7cpfnmdBXxzGPO3IXIQjm3XAidROiyOiD3AX1GE2m8oDpvrPfT6JMXh66+A1cC3SuNuAb4A3EERekcAPy71/SGwGXhE0uOVA0fED4CPUeyt9lEEzRl11lXpYop57aHYS/230uvsoThn+qcUh9oPAG9KzVdQfF5WUwTpUuBFqe0DFGH2BPBq4CeD1PBJ4Bjgt8BNwPJSDc+m138lxfnAHcC7S+3bgbsofrH8aAjztmEaOCFrZm1E0jXAroj4+1bXkgP/4adZm5HUCZwOHN3iUrLhQ2OzNiLp0xQXmS6PiF+1up5c+NDYzLLnPUIzy56D0Myy15YXS6ZOnRqdnZ2tLsPMxpiNGzc+HhHPuYFFWwZhZ2cnPT09rS7DzMYYSQ9XWz/ooXG6XdDadFulzZIuTOsvl3SfpHskrSjfYqii/0OSfpHeg+p0M7O2U885wn3ARRHRBRwLnC+pC1gDvCYijgR+CVy6nzHeFBFHRUR3wxWbmTXZoEEYEX0RcVda3gP0AjMiYnXp7hl3AjNHrkwzs5EzpKvG6S/ejwY2VDSdB9xco1sAq9PdfhcPuUIzsxFW98USSZMo3hC/JCJ2l9Z/lOLweVmNrsdHxE5JLwXWSLovIm6vMv5iiht9Mnv27MpmM7MRU9ceoaQJFCG4LCKWl9Yvorh9+nujxltUImJn+vgYsAKYW2O7qyOiOyK6Ozr87xnMbPTUc9VYFLcj6o2IK0rrFwAfBt4eEVX/r0O619rA/804kOK+d/c2o3Azs2apZ49wHsWdcuenP4HZJGkhcCXFzSfXpHVXAUiaLmlV6jsNWC/pboqbUd4UEbc0fxpmZsM36DnCiFhP9du4r6qyjojYRfE/MIiIbcBrGynQzGyk+b3GZpY9B6GZZc9BaGbZcxCaWfYchGaWPQehmWXPQWhm2XMQmln2HIRmlj0HoZllz0FoZtlzEJpZ9hyEZpY9B6GZZc9BaGbZcxCaWfYchGaWPQehmWXPQWhm2XMQmln2HIRmlj0HoZllz0FoZtlzEJpZ9hyEZpY9B6GZZW/QIJQ0S9JaSVskbZZ0YVp/uaT7JN0jaYWkg2v0XyDpfklbJV3S5PrNzBpWzx7hPuCiiOgCjgXOl9QFrAFeExFHAr8ELq3sKGkc8BXgFKALODP1NTNrG4MGYUT0RcRdaXkP0AvMiIjVEbEvbXYnMLNK97nA1ojYFhHPANcBpzandDOz5hjSOUJJncDRwIaKpvOAm6t0mQFsLz3fkdZVG3uxpB5JPf39/UMpy8ysIXUHoaRJwI3AkojYXVr/UYrD52WNFBIRV0dEd0R0d3R0NDKUmdmQjK9nI0kTKEJwWUQsL61fBLwNeHNERJWuO4FZpecz0zozs7ZRz1VjAUuB3oi4orR+AfBh4O0R8XSN7j8D5kh6maSJwBnAysbLNjNrnnoOjecBZwPzJW1Kj4XAlcBkYE1adxWApOmSVgGkiykXALdSXGS5PiI2j8REzMyGa9BD44hYD6hK06oa2+8CFpaer6q1rZlZO/A7S8wsew5CM8ueg9DMsucgNLPsOQjNLHsOQjPLnoPQzLLnIDSz7DkIzSx7DkIzy56D0Myy5yA0s+w5CM0sew5CM8ueg9DMsucgNLPsOQjNLHsOQjPLnoPQzLLnIDSz7DkIzSx7DkIzy56D0Myy5yA0s+w5CM0sew5CM8veoEEoaZaktZK2SNos6cK0/p3p+e8kde+n/0OSfiFpk6SeZhZvZtYM4+vYZh9wUUTcJWkysFHSGuBe4HTgX+oY400R8XgDdZqZjZhBgzAi+oC+tLxHUi8wIyLWAEga2QrNzEbYkM4RSuoEjgY2DKFbAKslbZS0eD9jL5bUI6mnv79/KGWZmTWk7iCUNAm4EVgSEbuH8BrHR8QxwCnA+ZLeUG2jiLg6Irojorujo2MIw5uZNaauIJQ0gSIEl0XE8qG8QETsTB8fA1YAc4dapJnZSKrnqrGApUBvRFwxlMElHZgusCDpQOBkiossZmZto549wnnA2cD89CcwmyQtlHSapB3AccBNkm4FkDRd0qrUdxqwXtLdwE+BmyLilhGYh5nZsNVz1Xg9UOvS8Ioq2+8CFqblbcBrGynQzGyk+Z0lZpY9B6GZZc9BaGbZcxCaWfYchGaWPQehmWXPQWhm2XMQmln2HIRmlj0HoZllz0FoZtlzEJpZ9hyEZpY9B6GZZc9BaGbZcxCaWfYchGaWPQehmWXPQWhm2XMQmln2HIRmlj0HoZllz0FoZtlzEJpZ9hyEZpa9QYNQ0ixJayVtkbRZ0oVp/TvT899J6t5P/wWS7pe0VdIlzSzezKwZ6tkj3AdcFBFdwLHA+ZK6gHuB04Hba3WUNA74CnAK0AWcmfqambWNQYMwIvoi4q60vAfoBWZERG9E3D9I97nA1ojYFhHPANcBpzZatJlZMw3pHKGkTuBoYEOdXWYA20vPd6R1ZmZto+4glDQJuBFYEhG7m12IpMWSeiT19Pf3N3t4M7Oa6gpCSRMoQnBZRCwfwvg7gVml5zPTuueIiKsjojsiujs6OobwEmZmjannqrGApUBvRFwxxPF/BsyR9DJJE4EzgJVDL9PMbOTUs0c4DzgbmC9pU3oslHSapB3AccBNkm4FkDRd0iqAiNgHXADcSnGR5fqI2DwiMzEzG6bxg20QEesB1WheUWX7XcDC0vNVwKrhFmhmNtL8zhIzy56D0Myy5yA0s+w5CM0sew5CM8ueg9DMsucgNLPsOQjNLHsOQjPLnoPQzLLnIDSz7DkIzSx7DkIzy56D0Myy5yA0s+w5CM0sew5CM8ueg9DMsucgNLPsOQjNLHsOQjPLnoPQzLLnIDSz7DkIzSx7DkIzy56D0MyyN2gQSpolaa2kLZI2S7owrX+JpDWSHkgfp9To/6ykTemxstkTMDNrVD17hPuAiyKiCzgWOF9SF3AJcFtEzAFuS8+r+a+IOCo93t6Uqs3MmmjQIIyIvoi4Ky3vAXqBGcCpwDfSZt8A3jFCNZqZjaghnSOU1AkcDWwApkVEX2p6BJhWo9sLJfVIulPSO/Yz9uK0XU9/f/9QyjIza0jdQShpEnAjsCQidpfbIiKAqNH1sIjoBt4D/JOkV1TbKCKujojuiOju6Oiotywzs4bVFYSSJlCE4LKIWJ5WPyrp0NR+KPBYtb4RsTN93Aaso9ijNDNrG/VcNRawFOiNiCtKTSuBc9LyOcC/V+k7RdIBaXkqMA/Y0mjRZmbNVM8e4TzgbGB+6c9gFgKfA94i6QHgpPQcSd2Svpb6Hg70SLobWAt8LiIchGbWVsYPtkFErAdUo/nNVbbvAd6fln8CHNFIgWZmI83vLDGz7DkIzSx7DkIzy56D0Myy5yA0s+w5CM0sew5CM8ueg9DMsucgNLPsOQjNLHsOQjPLnoPQzLLnIDSz7Km4uXR7kdQPPNzqOqqYCjze6iKaYKzMAzyXdtWuczksIp5zC/y2DMJ2Jakn/duB57WxMg/wXNrV820uPjQ2s+w5CM0sew7Cobm61QU0yViZB3gu7ep5NRefIzSz7HmP0Myy5yAEJC2QdL+krZIuqdJ+mKTbJN0jaZ2kmaW22ZJWS+qVtEVS56gWX6HBuXxe0uY0ly+nf+XaEpKukfSYpHtrtCvVuDXN5ZhS2zmSHkiPc6r1H03DnYukoyTdkb4m90h69+hWXrXWYX9dUvtBknZIunJ0Kq5TRGT9AMYBDwIvByYCdwNdFdt8FzgnLc8HvlVqWwe8JS1PAl78fJwL8CfAj9MY44A7gBNbOJc3AMcA99ZoXwjcTPEfFo8FNqT1LwG2pY9T0vKUFn+PDXcurwLmpOXpQB9w8PNxLqX2LwHXAle2ch6VD+8Rwlxga0Rsi4hngOuAUyu26QJ+mJbXDrRL6gLGR8QagIjYGxFPj07ZVQ17LkAAL6QI0AOACcCjI15xDRFxO/DkfjY5FfhmFO4EDpZ0KPBWYE1EPBkRvwHWAAtGvuLahjuXiPhlRDyQxtgFPAY854+BR1MDXxck/TEwDVg98pUOjYMQZgDbS893pHVldwOnp+XTgMmSDqH4jf2UpOWSfi7pcknjRrzi2oY9l4i4gyIY+9Lj1ojoHeF6G1FrrvV8DtrNoDVLmkvxS+rBUaxrOKrORdILgC8AF7ekqkE4COtzMfBGST8H3gjsBJ4FxgMnpPbXURySLmpRjfWqOhdJrwQOB2ZSfDPPl3RC68q0AWmP6lvAuRHxu1bXM0wfAlZFxI5WF1LN+FYX0AZ2ArNKz2emdb+XDktOB5A0CfiziHhK0g5gU0RsS23fozgvsnQU6q6mkbl8ALgzIvamtpuB44AfjUbhw1BrrjuBEyvWrxu1qoan5tdN0kHATcBH06Fmu6s1l+OAEyR9iOJc+kRJeyPiORf0WsF7hPAzYI6kl0maCJwBrCxvIGlq2rUHuBS4ptT3YEkD523mA1tGoeZaGpnLryn2FMdLmkCxt9jOh8Yrgfelq5THAr+NiD7gVuBkSVMkTQFOTuvaWdW5pK/hCopzbje0tsS6VZ1LRLw3ImZHRCfFUck32yUEwXuERMQ+SRdQ/LCMA66JiM2SPgX0RMRKij2Mz0oK4Hbg/NT3WUkXA7elPzXZCHy1FfNI9Qx7LsANFEH+C4oLJ7dExH+M9hwGSPoORa1T0573xyku4BARVwGrKK5QbgWeBs5NbU9K+jTFLwWAT0XE/k7uj7jhzgV4F8VV2kMkLUrrFkXEptGqvVIDc2lrfmeJmWXPh8Zmlj0HoZllz0FoZtlzEJpZ9hyEZpY9B6GZZc9BaGbZcxCaWfb+F+T2uAYYHBdbAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 360x216 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\"--data_path\", default=\"../data\", type=str, help=\"The input data dir\")\n",
    "parser.add_argument(\"--batch_size\", default=128, type=int, help=\"The batch size of training\")\n",
    "parser.add_argument(\"--device\", default='cuda', type=str, help=\"The training device\")\n",
    "parser.add_argument(\"--learning_rate\", default=0.0004, type=float, help=\"learning rate\")\n",
    "parser.add_argument(\"--epochs\", default=1, type=int, help=\"Training epoch\")\n",
    "parser.add_argument(\"--modeldir\", default=\"./model\", type=str)\n",
    "args = parser.parse_known_args()[0]\n",
    "\n",
    "train_loader, test_loader = cifar100_dataset(args)\n",
    "\n",
    "writer = SummaryWriter(os.path.join(args.modeldir, \"tensorboard\"))\n",
    "net = SE_ResNet18().to(args.device)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=args.learning_rate)\n",
    "\n",
    "lossv, accv = [], []\n",
    "index_num = 0\n",
    "for epoch in range(args.epochs):\n",
    "    train(net, train_loader, optimizer, criterion, writer, args, epoch, index_num)\n",
    "    PATH = os.path.join(args.modeldir, 'SE_ResNet18.pth')\n",
    "    torch.save(net.state_dict(), PATH)\n",
    "    with torch.no_grad():\n",
    "        validate(net, test_loader, criterion, writer, args, epoch, lossv, accv)\n",
    "plt.figure(figsize=(5, 3))\n",
    "plt.plot(np.arange(1, args.epochs + 1), lossv)\n",
    "plt.title('validation loss')\n",
    "plt.savefig(os.path.join(args.modeldir, 'validation_loss'))\n",
    "\n",
    "plt.figure(figsize=(5, 3))\n",
    "plt.plot(np.arange(1, args.epochs + 1), accv)\n",
    "plt.title('validation accuracy')\n",
    "plt.savefig(os.path.join(args.modeldir, 'validation_accuracy'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2f394aca7ca06fed1e6064aef884364492d7cdda3614a461e02e6407fc40ba69"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
